package org.example.zk_util

import org.apache.zookeeper.*
import org.apache.zookeeper.data.Stat
import java.util.concurrent.CountDownLatch
import java.util.concurrent.ConcurrentHashMap

class ZkClient {
    private lateinit var zk: ZooKeeper
    private val nodeDataListeners = ConcurrentHashMap<String, (String, String) -> Unit>() // path, (path, data)
    private val nodeChildrenListeners = ConcurrentHashMap<String, (String, List<String>) -> Unit>() // path, (path, children)

    fun connect(connectString: String) {
        System.setProperty("java.security.auth.login.config", "zk_client_jaas.conf")

        val latch = CountDownLatch(1)
        zk = ZooKeeper(connectString, 5000, {
            if (it.state == Watcher.Event.KeeperState.SyncConnected) {
                latch.countDown()
            }
        })
        latch.await()
        println("连接成功")
        // 添加认证信息 (根据需要取消注释并修改)
        // val username = "admin"
        // val password = "admin123"
        // zk.addAuthInfo("digest", "$username:$password".toByteArray())
    }

    fun close() {
        zk.close()
    }

    fun getNodeData(path: String): String {
        val stat = Stat()
        val data = zk.getData(path, false, stat)
        return if (data.isEmpty()) "" else String(data)
    }

    fun setNodeData(path: String, data: String): Boolean {
        try {
            zk.setData(path, data.toByteArray(), -1)
            return true
        } catch (e: Exception) {
            e.printStackTrace()
            return false
        }
    }

    fun createNode(path: String, data: String = ""): Boolean {
        try {
            zk.create(path, data.toByteArray(), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT)
            return true
        } catch (e: Exception) {
            e.printStackTrace()
            return false
        }
    }

    fun deleteNode(path: String): Boolean {
        try {
            zk.delete(path, -1)
            return true
        } catch (e: Exception) {
            e.printStackTrace()
            return false
        }
    }

    fun getChildren(path: String): List<String> {
        return zk.getChildren(path, false)
    }

    fun watchNode(path: String, listener: (String, String) -> Unit) {
        nodeDataListeners[path] = listener
        zk.getData(path, {
            if (it.type == Watcher.Event.EventType.NodeDataChanged) {
                val data = getNodeData(path)
                listener(path, data)
                // 重新注册监听
                watchNode(path, listener)
            }
        }, null)
    }

    fun watchChildren(path: String, listener: (String, List<String>) -> Unit) {
        nodeChildrenListeners[path] = listener
        zk.getChildren(path, {
            if (it.type == Watcher.Event.EventType.NodeChildrenChanged) {
                val children = getChildren(path)
                listener(path, children)
                // 重新注册监听
                watchChildren(path, listener)
            }
        }, null)
    }

    fun traverseNode(path: String): Map<String, Any> {
        val result = mutableMapOf<String, Any>()
        val stat = Stat()
        val data = zk.getData(path, false, stat)
        result["data"] = if (data.isEmpty()) "" else String(data)
        result["version"] = stat.version
        result["numChildren"] = stat.numChildren

        val children = zk.getChildren(path, false)
        if (children.isNotEmpty()) {
            val childrenMap = mutableMapOf<String, Any>()
            for (child in children) {
                val childPath = if (path == "/") "/$child" else "$path/$child"
                childrenMap[child] = traverseNode(childPath)
            }
            result["children"] = childrenMap
        }
        return result
    }
}